#!/usr/bin/env python2.7
#-*- coding=utf-8 -*-

from .cstructures import CStruct
from .protocols import (
        LinkLayer,
        InternetLayer,
        TransportLayer,
        ApplicationLayer
    )
import struct

class _PcapMagic_(CStruct):
    def __init__(self, data):
        super(_PcapMagic_, self).__init__(data)
        magic = getattr(self, 'magic', None)
        if magic == '\xd4\xc3\xb2\xa1':
            self.endian = '<'
        elif magic == '\xa1\xb2\xc3\xd4':
            self.endian = '>'
        else:
            raise ValueError('File has no valid MagicNumber[%s]' % magic)

    _fields_ = [('magic', 'CString', 4)]

class PcapHeader(CStruct):
    def __init__(self, data):
        magic = _PcapMagic_(data)
        self.magic = magic.magic
        self.endian = magic.endian
        super(PcapHeader, self).__init__(data, self.endian)
        self.datalen += 4

    def getnames(self):
        names = ['magic']
        names.extend(super(PcapHeader, self).getnames())
        return names

    _fields_ = [
            ('vermaj', 'WORD'),
            ('vermin', 'WORD'),
            ('thiszone', 'Int32'),
            ('sigfigs', 'DWORD'),
            ('snaplen', 'DWORD'),
            ('network', 'DWORD')
        ]

class PacketHeader(CStruct):
    _fields_ = [
            ('ts_sec', 'DWORD'),
            ('ts_usec', 'DWORD'),
            ('incl_len', 'DWORD'),
            ('orig_len', 'DWORD'),
        ]

class PcapReader(object):
    def __init__(self, fpath):
        self.fp = open(fpath, 'rb')
        pcaphdr = PcapHeader(self.fp)
        self._endian_ = pcaphdr.endian
        self.network = pcaphdr.network

    def readpacket(self):
        try:
            packhdr = PacketHeader(self.fp, self._endian_)
        except BufferError:
            return None
        datalen = packhdr.incl_len
        packdata = self.fp.read(datalen)
        return packdata

    def seekpacket(self, num):
        if type(num) is not int or num < 1:
            raise ValueError('"num" must be int type and not less than 1!')
        self.fp.seek(0x20)
        count = 1
        while count <= num:
            offsetdata = self.fp.read(4)
            if len(offsetdata) < 4:
                return None
            offset = struct.unpack('%sI' % self._endian_, offsetdata)[0]
            if count == num:
                self.fp.seek(0x4, 1)
                data = self.fp.read(offset)
                if len(data) == offset:
                    return data
            else:
                self.fp.seek(offset + 0xc, 1)
            count += 1
        return None

    def __enter__(self):
        return self

    def __exit__(self, *exc_info):
        self.close()

    def close(self):
        self.fp.close()

class PcapReaderEx(PcapReader):
    def _resolvelayer_(self, layer, layertype, layerdata):
        try:
            layerclass = layer.layerdict[layertype]
            thislayer = layerclass(layerdata)
            nextlayer = getattr(thislayer, 'nextlayer', None)
            restdata = getattr(thislayer, 'restdata', None)
            classname = getattr(thislayer, 'classname', '')
        except (KeyError, BufferError):
            return None
        return thislayer, restdata, nextlayer

    def _resolvepacket_(self, packdata):
        ret = {
                'LinkLayer':None,
                'InternetLayer':None,
                'TransportLayer':None,
                'ApplicationLayer':None,
                'Data':packdata
            }
        layertype = self.network
        for layer in ((LinkLayer, 'LinkLayer'), (InternetLayer, 'InternetLayer'), (TransportLayer, 'TransportLayer')):
            layerclass = layer[0]
            layername = layer[1]
            result = self._resolvelayer_(layerclass, layertype, packdata)
            if result is None:
                break
            ret[layername] = result[0]
            packdata = result[1]
            ret['Data'] = packdata
            layertype = result[2]
        return ret

    def readpacket(self):
        packdata = super(PcapReaderEx, self).readpacket()
        if packdata is None:
            return None
        return self._resolvepacket_(packdata)

    def seekpacket(self, num):
        packdata = super(PcapReaderEx, self).seekpacket(num)
        if packdata is None:
            return None
        return self._resolvepacket_(packdata)
